#ifndef COST_FUNCTION_H
#define COST_FUNCTION_H
#include <iostream>
#include <vector>

#include <ceres/ceres.h>

#include "math_utils.h"
#include "state.h"

struct PoseNumericalDiff {
    PoseNumericalDiff(const Eigen::Vector3d &last_pose, const Eigen::Vector2d &last_velocity, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity) 
        : last_pose_(last_pose), last_velocity_(last_velocity), pose_(pose), velocity_(velocity) {}

    bool operator()(const double* const state, double* res) const {
        Eigen::Matrix<double, 13, 1> state_tmp;
        for(int i = 0; i < 13; i++) {
            state_tmp[i] = state[i];
        }
        Eigen::Matrix<double, 10, 1> residual;
        Eigen::Vector2d res_p = state_tmp.head<2>() - last_pose_.block<2,1>(0,0);
        double res_theta = math_utils::theta_normalize(state_tmp(4,0) - last_pose_[2]);
        Eigen::Vector2d res_v = state_tmp.block<2,1>(2,0) - last_velocity_;

        (residual).block<2,1>(0,0) = res_p;
        (residual).block<2,1>(2,0) = res_v;
        (residual)(4,0) = res_theta; 

        (residual).block<2,1>(5,0) = state_tmp.block<2,1>(8,0) - pose_.block<2,1>(0,0);
        (residual).block<2,1>(7,0) = state_tmp.block<2,1>(10,0) - velocity_.block<2,1>(0,0);
        (residual)(9,0) = math_utils::theta_normalize(state_tmp(12,0) - pose_[2]);  
        
        for(int i = 0; i < 10; i++) {
            res[i] = residual[i];
        }
        return true;
    }

private:
    Eigen::Vector3d last_pose_;
    Eigen::Vector2d last_velocity_;
    Eigen::Vector3d pose_;
    Eigen::Vector2d velocity_;    
};

struct PoseAutoDiff {
    PoseAutoDiff(const Eigen::Vector3d &last_pose, const Eigen::Vector2d &last_velocity, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity) 
        : last_pose_(last_pose), last_velocity_(last_velocity), pose_(pose), velocity_(velocity) {}

    template <typename T>    
    bool operator()(const T* const state, T* res) const {
        Eigen::Matrix<T, 13, 1> state_tmp;
        for(int i = 0; i < 13; i++) {
            state_tmp[i] = (state[i]);
        }
        Eigen::Matrix<T, 10, 1> residual;
        // Eigen::Vector2d res_p = state_tmp.head<2>() - last_pose_.block<2,1>(0,0);
        // double res_theta = math_utils::theta_normalize(state_tmp(4,0) - last_pose_[2]);
        // Eigen::Vector2d res_v = state_tmp.block<2,1>(2,0) - last_velocity_;

        // (residual).block<2,1>(0,0) = res_p;
        // (residual).block<2,1>(2,0) = res_v;
        // (residual)(4,0) = res_theta; 

        // (residual).block<2,1>(5,0) = state_tmp.block<2,1>(8,0) - pose_.block<2,1>(0,0);
        // (residual).block<2,1>(7,0) = state_tmp.block<2,1>(10,0) - velocity_.block<2,1>(0,0);
        // (residual)(9,0) = math_utils::theta_normalize(state_tmp(12,0) - pose_[2]);  
        
        for(int i = 0; i < 10; i++) {
            res[i] = T(residual[i]);
        }
        return true;
    }

private:
    Eigen::Vector3d last_pose_;
    Eigen::Vector2d last_velocity_;
    Eigen::Vector3d pose_;
    Eigen::Vector2d velocity_;    
};

struct PoseAnalyticDiff  : public ceres::CostFunction {
    
    PoseAnalyticDiff(const Eigen::Vector3d &last_pose, const Eigen::Vector2d &last_velocity, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity) 
        : last_pose_(last_pose), last_velocity_(last_velocity), pose_(pose), velocity_(velocity) {
            
    
    }

    bool operator()(const double* const state, double* res) const {
        Eigen::Matrix<double, 13, 1> state_tmp;
        for(int i = 0; i < 13; i++) {
            state_tmp[i] = state[i];
        }
        Eigen::Matrix<double, 10, 1> residual;
        Eigen::Vector2d res_p = state_tmp.head<2>() - last_pose_.block<2,1>(0,0);
        double res_theta = math_utils::theta_normalize(state_tmp(4,0) - last_pose_[2]);
        Eigen::Vector2d res_v = state_tmp.block<2,1>(2,0) - last_velocity_;

        (residual).block<2,1>(0,0) = res_p;
        (residual).block<2,1>(2,0) = res_v;
        (residual)(4,0) = res_theta; 

        (residual).block<2,1>(5,0) = state_tmp.block<2,1>(8,0) - pose_.block<2,1>(0,0);
        (residual).block<2,1>(7,0) = state_tmp.block<2,1>(10,0) - velocity_.block<2,1>(0,0);
        (residual)(9,0) = math_utils::theta_normalize(state_tmp(12,0) - pose_[2]);  
        
        for(int i = 0; i < 10; i++) {
            res[i] = residual[i];
        }
        return true;
    }

private:
    Eigen::Vector3d last_pose_;
    Eigen::Vector2d last_velocity_;
    Eigen::Vector3d pose_;
    Eigen::Vector2d velocity_;    
};

class CeresCostFunction : public ceres::CostFunction {
public:
    CeresCostFunction() {
        // 设置残差维度(10)和参数块维度(13)
        set_num_residuals(10);
        mutable_parameter_block_sizes()->push_back(13);
    }

    virtual bool Evaluate(double const* const* parameters,
                        double* residuals,
                        double** jacobians) const override {
        const double* params = parameters[0];  // 获取13维参数

        // 1. 计算残差 (10维)
        for (int i = 0; i < 10; ++i) {
            residuals[i] = 0.0;
            for (int j = 0; j < 13; ++j) {
                residuals[i] += (j + 1) * params[j] * (i + 1);  // 示例计算
            }
            residuals[i] -= (i + 1) * 10;  // 目标值偏移
        }

        // 2. 计算雅可比矩阵 (10x13)
        if (jacobians != nullptr && jacobians[0] != nullptr) {
            double* jacobian = jacobians[0];  // 指向10x13矩阵的指针
            
            // 行优先存储: 第i个残差对第j个参数的导数在jacobian[i*13 + j]
            for (int i = 0; i < 10; ++i) {
                for (int j = 0; j < 13; ++j) {
                    jacobian[i * 13 + j] = (j + 1) * (i + 1);  // ∂residual_i/∂param_j
                }
            }
        }

        return true;
    }
};

struct ExponentialResidual {
    ExponentialResidual(double x, double y) : x_(x), y_(y) {}
    
    template <typename T>
    bool operator()(const T* const m, const T* const c, T* residual) const {
        residual[0] = y_ - exp(m[0] * T(x_) + c[0]);
        return true;
    }
    
private:
    const double x_;
    const double y_;
};

#endif 
